"""
Instance Sequence Matching
Modified from DETR (https://github.com/facebookresearch/detr)
"""
import torch
from scipy.optimize import linear_sum_assignment
from torch import nn
import torch.nn.functional as F
from .point_feature import get_uncertain_point_coords_with_randomness,point_sample,calculate_uncertainty
from util.box_ops import box_cxcywh_to_xyxy, generalized_box_iou, multi_iou
from util.misc import nested_tensor_from_tensor_list
from torch.cuda.amp import autocast
INF = 100000000

def batch_dice_loss(inputs: torch.Tensor, targets: torch.Tensor):
    """
    Compute the DICE loss, similar to generalized IOU for masks
    Args:
        inputs: A float tensor of arbitrary shape.
                The predictions for each example.
        targets: A float tensor with the same shape as inputs. Stores the binary
                 classification label for each element in inputs
                (0 for the negative class and 1 for the positive class).
    """
    inputs = inputs.sigmoid()
    inputs = inputs.flatten(1)
    numerator = 2 * torch.einsum("nc,mc->nm", inputs, targets)
    denominator = inputs.sum(-1)[:, None] + targets.sum(-1)[None, :]
    loss = 1 - (numerator + 1) / (denominator + 1)
    return loss


batch_dice_loss_jit = torch.jit.script(
    batch_dice_loss
)  # type: torch.jit.ScriptModule


def batch_sigmoid_ce_loss(inputs: torch.Tensor, targets: torch.Tensor):
    """
    Args:
        inputs: A float tensor of arbitrary shape.
                The predictions for each example.
        targets: A float tensor with the same shape as inputs. Stores the binary
                 classification label for each element in inputs
                (0 for the negative class and 1 for the positive class).
    Returns:
        Loss tensor
    """
    hw = inputs.shape[1]

    pos = F.binary_cross_entropy_with_logits(
        inputs, torch.ones_like(inputs), reduction="none"
    )
    neg = F.binary_cross_entropy_with_logits(
        inputs, torch.zeros_like(inputs), reduction="none"
    )

    loss = torch.einsum("nc,mc->nm", pos, targets) + torch.einsum(
        "nc,mc->nm", neg, (1 - targets)
    )

    return loss / hw


# batch_sigmoid_ce_loss_jit = torch.jit.script(
#     batch_sigmoid_ce_loss
# )  # type: torch.jit.ScriptModule





def dice_coef(inputs, targets):
    inputs = inputs.sigmoid()
    inputs = inputs.flatten(1).unsqueeze(1) # [N, 1, THW]
    targets = targets.flatten(1).unsqueeze(0) # [1, M, THW]
    numerator = 2 * (inputs * targets).sum(2)
    denominator = inputs.sum(-1) + targets.sum(-1)

    # NOTE coef doesn't be subtracted to 1 as it is not necessary for computing costs
    coef = (numerator + 1) / (denominator + 1)
    return coef

def sigmoid_focal_coef(inputs, targets, alpha: float = 0.25, gamma: float = 2):
    N, M = len(inputs), len(targets)
    inputs = inputs.flatten(1).unsqueeze(1).expand(-1, M, -1) # [N, M, THW]
    targets = targets.flatten(1).unsqueeze(0).expand(N, -1, -1) # [N, M, THW]

    prob = inputs.sigmoid()
    ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
    p_t = prob * targets + (1 - prob) * (1 - targets)
    coef = ce_loss * ((1 - p_t) ** gamma)

    if alpha >= 0:
        alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
        coef = alpha_t * coef

    return coef.mean(2) # [N, M]


class NoBoxHungarianMatcher(nn.Module):
    """This class computes an assignment between the targets and the predictions of the network

    For efficiency reasons, the targets don't include the no_object. Because of this, in general,
    there are more predictions than targets. In this case, we do a 1-to-1 matching of the best predictions,
    while the others are un-matched (and thus treated as non-objects).
    """

    def __init__(self, cost_class: float = 1,
                       cost_mask: float = 1, cost_dice: float = 1, num_classes: int = 1):
        """Creates the matcher

        Params:
            cost_class: This is the relative weight of the classification error in the matching cost
            cost_bbox: This is the relative weight of the L1 error of the bounding box coordinates in the matching cost
            cost_giou: This is the relative weight of the giou loss of the bounding box in the matching cost
            cost_mask: This is the relative weight of the sigmoid focal loss of the mask in the matching cost
            cost_dice: This is the relative weight of the dice loss of the mask in the matching cost
        """
        super().__init__()
        self.cost_class = cost_class
        self.cost_mask = cost_mask
        self.cost_dice = cost_dice
        self.num_classes = num_classes
        assert cost_class != 0 \
            or cost_mask != 0 or cost_dice != 0, "all costs cant be 0"
        self.mask_out_stride = 4

    @torch.no_grad()
    def forward(self, outputs, targets):
        """ Performs the matching
        Params:
            outputs: This is a dict that contains at least these entries:
                 "pred_logits": Tensor of dim [batch_size, num_queries_per_frame, num_frames, num_classes] with the classification logits
                 "pred_boxes": Tensor of dim [batch_size, num_queries_per_frame, num_frames, 4] with the predicted box coordinates
                 "pred_masks": Tensor of dim [batch_size, num_queries_per_frame, num_frames, h, w], h,w in 4x size
            targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing:
                 NOTE: Since every frame has one object at most
                 "labels": Tensor of dim [num_frames] (where num_target_boxes is the number of ground-truth
                           objects in the target) containing the class labels
                 "boxes": Tensor of dim [num_frames, 4] containing the target box coordinates
                 "masks": Tensor of dim [num_frames, h, w], h,w in origin size 
        Returns:
            A list of size batch_size, containing tuples of (index_i, index_j) where:
                - index_i is the indices of the selected predictions (in order)
                - index_j is the indices of the corresponding selected targets (in order)
            For each batch element, it holds:
                len(index_i) = len(index_j) = min(num_queries, num_target_boxes)
        """
        src_logits = outputs["pred_logits"] 
        src_masks = outputs["pred_masks"]   

        nf, nq, h, w = src_masks.shape

        masks = [t for t in targets[0]['masks']]

        # handle mask padding issue
        # target_masks, valid = nested_tensor_from_tensor_list(masks, 
        #                                                      size_divisibility=32,
        #                                                      split=False).decompose()
        target_masks, valid = nested_tensor_from_tensor_list(masks).decompose()       
        target_masks = target_masks.to(src_masks) # [B, T, H, W]

        indices = []   
        tgt_valid = targets[0]["valid"]    

        for t in range(nf):
            out_prob_split = src_logits[t].softmax(-1)    
            out_mask_split = src_masks[t]
            tgt_mask_split = target_masks[t] #[n,h,w]
            num_queries = out_mask_split.size(0)


            valid_count = (tgt_valid[t] == 1).sum().item()

            # Compute the classification cost.
            tgt_ids_split = torch.tensor([0] * valid_count).long()
            cost_class = -out_prob_split[:, tgt_ids_split]

            tgt_mask_split = tgt_mask_split[tgt_valid[t] == 1]

            out_mask = out_mask_split[:, None]
            tgt_mask = tgt_mask_split[:, None]

            t,h,w = out_mask.size()[-3:]
            tgt_mask = F.interpolate(tgt_mask, size=(h, w), mode='bilinear', align_corners=False) 
            out_mask = out_mask.reshape(out_mask.size()[0],t*h*w)
            tgt_mask = tgt_mask.reshape(tgt_mask.size()[0],t*h*w)

            # point_coords = torch.rand(1, 12544, 2, device=out_mask.device)

            # tgt_mask = point_sample(
            #     tgt_mask,
            #     point_coords.repeat(tgt_mask.shape[0], 1, 1),
            #     align_corners=False,
            # ).squeeze(1) #[K,12544]

            # out_mask = point_sample(
            #     out_mask,
            #     point_coords.repeat(out_mask.shape[0], 1, 1),
            #     align_corners=False,
            # ).squeeze(1)

            with autocast(enabled=False):
                out_mask = out_mask.float()
                tgt_mask = tgt_mask.float()

                cost_mask = batch_sigmoid_ce_loss(out_mask, tgt_mask) #[N,K]
                cost_dice = batch_dice_loss(out_mask, tgt_mask)

        
            # Final cost matrix
            C = self.cost_class * cost_class + \
                self.cost_mask * cost_mask + self.cost_dice * cost_dice  # [q, 1]
        
            C = C.reshape(num_queries, -1).cpu()
            indices.append(linear_sum_assignment(C))

        return [
            (torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64))
            for i, j in indices
        ]
                




def build_nobox_matcher(args):
    num_classes = 1
    # if args.binary:
    #     num_classes = 1
    # else:
    #     if args.dataset_file == 'ytvos':
    #         num_classes = 65 
    #     elif args.dataset_file == 'davis':
    #         num_classes = 78
    #     elif args.dataset_file == 'a2d' or args.dataset_file == 'jhmdb':
    #         num_classes = 1
    #     else: 
    #         num_classes = 91  # for coco
    return NoBoxHungarianMatcher(cost_class=args.set_cost_class, 
                            #cost_bbox=args.set_cost_bbox, 
                            #cost_giou=args.set_cost_giou,
                            cost_mask=args.set_cost_mask,
                            cost_dice=args.set_cost_dice,
                            num_classes=num_classes)